import time
import paramiko
import re
from robot.api.deco import keyword

@keyword("Try Connect SSH")
def try_connect_SSH(hostname, port, username, password):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        # 尝试建立SSH连接
        ssh.connect(hostname, int(port), username, password, timeout=10)
        print("✅ ssh连接成功!")
        return ssh
    except Exception as e:
        print(f"❌ ssh连接失败: {str(e)}")
        return None

@keyword("Send Ssh Msg")
def send_ssh_msg(ssh, command, timeout=60):
    if not ssh:
        return "❌ SSH连接未建立!"
    
    try:
        shell = ssh.invoke_shell()
        shell.settimeout(timeout)
        
        # 等待shell准备就绪
        time.sleep(2)
        
        # 清空缓冲区
        while shell.recv_ready():
            shell.recv(1024)
        
        # 分块发送命令，确保长命令完整发送
        command = command + "\n"
        chunk_size = 100  # 每次发送100个字符
        
        for i in range(0, len(command), chunk_size):
            chunk = command[i:i+chunk_size]
            shell.send(chunk)
            time.sleep(0.5)  # 小延迟确保发送完整
        
        # 等待命令执行
        time.sleep(3)
        
        output = ""
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            if shell.recv_ready():
                # 增加缓冲区大小
                chunk = shell.recv(65536).decode('utf-8', 'ignore')
                output += chunk
                
                # 检查是否回到命令提示符
                if re.search(r'[$#]\s*$|~.*[$#]\s*$', output.split('\n')[-1]):
                    # 额外等待确保输出完整
                    time.sleep(2)
                    if not shell.recv_ready():
                        break
            else:
                time.sleep(0.5)
                # 如果10秒没有新数据，认为完成
                if time.time() - start_time > 10 and not shell.recv_ready():
                    break
        
        return output
        
    except Exception as e:
        return f"❌ 命令执行错误: {str(e)}"

@keyword("Close SSH Connection")
def close_ssh_connection(ssh):
    if ssh:
        ssh.close()
        print("✅ SSH连接已关闭!")